package com.example.ai.etl.reader;

import net.sourceforge.tess4j.ITessAPI;
import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.ParagraphPdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

@Component
public class MyPagePdfDocumentReader {

    public List<Document> getDocsFromPdfWithCatalog(String classPathFile) {
        ParagraphPdfDocumentReader pdfReader = new ParagraphPdfDocumentReader("classpath:/" + classPathFile,
                PdfDocumentReaderConfig.builder()
                        .withPageTopMargin(0)
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfTopTextLinesToDelete(0)
                                .build())
                        .withPagesPerDocument(1)
                        .build());

        return pdfReader.read();
    }

    public List<Document> getDocs(String classPathFile) {
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader("classpath:/" + classPathFile,
                PdfDocumentReaderConfig.builder()
                        .withPageTopMargin(0)
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfTopTextLinesToDelete(0)
                                .build())
                        .withPagesPerDocument(1)
                        .build());
        return pdfReader.read();
    }

    // 新增方法：处理文字不可选择的PDF文件
    List<Document> getDocsFromUnselectablePdf(String classPathFile) {
        try {
            // 使用Tesseract OCR工具提取文本
            ITesseract tesseract = new Tesseract();
            tesseract.setDatapath("tessdata"); // 设置Tesseract数据路径
            tesseract.setLanguage("chi_sim"); // 设置语言为简体中文

            tesseract.setOcrEngineMode(ITessAPI.TessOcrEngineMode.OEM_LSTM_ONLY);
            tesseract.setPageSegMode(ITessAPI.TessPageSegMode.PSM_SINGLE_BLOCK);
            // 将ClassPathResource转换为File对象
            ClassPathResource resource = new ClassPathResource(classPathFile);
            File pdfFile = resource.getFile();

            // 提取PDF文本
            String extractedText = tesseract.doOCR(pdfFile);

            // 将提取的文本封装为Document对象
            List<Document> documents = new ArrayList<>();
            Document document = new Document(extractedText);
            documents.add(document);

            return documents;
        } catch (IOException | TesseractException e) {
            throw new RuntimeException("Failed to extract text from unselectable PDF", e);
        }
    }

    public static void main(String[] args) {
        MyPagePdfDocumentReader myPagePdfDocumentReader = new MyPagePdfDocumentReader();
        /*
        List<Document> documents = myPagePdfDocumentReader.getDocsFromPdfWithCatalog("中华人民共和国劳动法.pdf");
        System.out.println(documents);*/

        // 测试处理文字不可选择的PDF文件
        List<Document> documents = myPagePdfDocumentReader.getDocsFromUnselectablePdf("data/劳动合同.pdf");
        System.out.println(documents);
    }
}